package com.wf.interceptor;


import cn.hutool.core.io.resource.ResourceUtil;

import com.wf.exception.BankException;
import com.wf.interceptor.annotation.RateLimit;
import com.wf.utils.ServletUtil;
import org.springframework.context.annotation.Lazy;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;

/**
 * 限 流 拦 截 器
 * <p>
 * Author: 就 眠 仪 式
 * CreateTime: 2020/10/23
 */
@Component
public class RateLimitInterceptor implements HandlerInterceptor {

    @Lazy
    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            RateLimit currentLimit = handlerMethod.getMethodAnnotation(RateLimit.class);
            // 存在注解
            if (currentLimit != null) {
                // 限流条件
                int number = currentLimit.number();
                long time = currentLimit.time();
                if (time > 0 && number > 0) {
                    String key = request.getContextPath() + ":" + request.getServletPath() + ":" + ServletUtil.getIpAddress();

                    List<String> keys = Collections.singletonList(key);

                    DefaultRedisScript<Boolean> script = new DefaultRedisScript<>();

                    script.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/rateLimit.lua")));

                    script.setResultType(Boolean.class);

                    Boolean result = redisTemplate.execute(script, keys, number, time);

                    if (!result) {
                        throw new BankException("请求频繁，请稍后重试！");
                    }
                    return true;
                }
            }
        }
        return true;
    }
}
